import os
import sqlite3

import dotenv
from typing import TypedDict, Annotated, List
from langchain_core.messages import BaseMessage, AIMessage, ToolMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.sqlite import SqliteSaver

# 从本地文件导入工具
from tools import available_tools

dotenv.load_dotenv()


# 1. 定义Graph的状态
class AgentState(TypedDict):
    messages: Annotated[List[BaseMessage], lambda x, y: x + y]


# 2. 定义节点和逻辑
# LangGraph已经提供了封装好的ToolNode，可以直接使用
tool_node = ToolNode(available_tools)
model = ChatOpenAI(
    model="deepseek-chat",
    base_url=os.getenv("DS_BASE"),
    api_key=os.getenv("DS_API_KEY"),
    temperature=0
).bind_tools(available_tools)


def agent_node(state: AgentState):
    """Agent决策节点：调用LLM决定下一步行动。"""
    print("---NODE: Agent is thinking...---")

    # 确保只将必要的消息发送给LLM
    # 如果最后一条消息是ToolMessage，需要让LLM重新规划
    last_message = state["messages"][-1]
    if isinstance(last_message, ToolMessage):
        print("---Agent received feedback. Re-planning...---")

    response = model.invoke(state["messages"])
    print(f"AI thinking ret: {response}")
    return {"messages": [response]}


def should_continue(state: AgentState) -> str:
    """条件路由：判断是否需要调用工具。"""
    print("---ROUTER: Deciding next step---")
    last_message = state["messages"][-1]

    if last_message.tool_calls:
        print("Go continue (call tools)")
        return "continue"

    print("Go end")
    return "end"


# 3. 构建并编译Graph
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
# 【核心修复】直接使用 LangGraph 提供的 ToolNode 实例
workflow.add_node("tools", tool_node)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
    "agent",
    should_continue,
    {"continue": "tools", "end": END},
)
workflow.add_edge("tools", "agent")

# 设置内存/检查点，这是实现多轮交互和暂停的关键
conn = sqlite3.connect("checkpoints.sqlite",
                       check_same_thread=False,
                       isolation_level=None,
                       timeout=30,  # 等待锁的超时时间
                       detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES)
memory = SqliteSaver(conn)

# 编译Graph，并设置中断条件
app_graph = workflow.compile(
    checkpointer=memory,
    interrupt_before=["tools"]  # <-- 实现HITL的关键
)
